import pymongo ;
import os,collections ;
import re ;
import datetime;
from pymongo import MongoClient;
import time
from KDataBase import KGlobals
from KDataBase.KDBUtility import KDBUtility
from PyQt4.QtCore import QObject
import numpy as np
import pandas as pd
class KQueryDB(QObject):
	DbClient=MongoClient()
	Db=DbClient[KGlobals.DatabaseName]

	@staticmethod
	def QueryAllCountryMarkets():
		result=KQueryDB.__QueryAllStockCollections()
		countryset=set()
		marketset=set()
		for each_collection in result:
			(country,market)=KDBUtility.GetMarketCountryByCollectionName(each_collection)
			countryset.add(country)
			marketset.add(market)
		return list(countryset),list(marketset)	
	@staticmethod
	def __QueryAllStockCollections():
	# query all stock collections in the database
		results=[]
		allcollection=KQueryDB.Db.collection_names()
		for each_collection in allcollection:
			if not KDBUtility.IsSymbolBased(each_collection) and each_collection !='system.indexes':
				results.append(each_collection)
		return results
	@staticmethod
	def QueryAllStockNames(FullName=True):
	# Query All Stock names in the whole database
		allcollections=KQueryDB.__QueryAllStockCollections()
		results=[]
		for each_collection in allcollections:
			(country,market)=each_collection.split(".",2)
			stocknames=KQueryDB.QueryAllStocksByMkt(country,market,FullName)
			results.extend(stocknames)
		return results
	@staticmethod
	def QueryStockFullName(shortname,targetcountry):
		allcollections=KQueryDB.__QueryAllStockCollections()
		results=[]
		for each_collection in allcollections:
			(country,market)=each_collection.split(".",2)
			if targetcountry!=country:
				continue
			cldailystocks=KQueryDB.Db[KDBUtility.GetCollectionName(country,market,True)]
			cursor=cldailystocks.find_one({"symbol":shortname.upper()},{"symbol":1})
			if cursor:
				return "%s.%s.%s" %(shortname,market,country)
		return ''
	@staticmethod
	def QueryAllStocksByMkt(country,market,FullName=False,WriteToFile=''):
		# get the symbol name list of all stocks in this market
		cldailystocks=KQueryDB.Db[KDBUtility.GetCollectionName(country,market,True)]
		allsymobls=cldailystocks.find({},{"symbol":1}).distinct("symbol")
		if FullName==True:
			allsymobls=["%s.%s.%s" %(x,market,country) for x in allsymobls]
		if WriteToFile!='':
			file = open(WriteToFile,'w+')
			for item in allsymobls:
				file.write("%s\n" % item)
		return allsymobls	
	
	@staticmethod
	def QueryAllStocksStaticInfoByMkt(tcountry,tmarket):		
		allcollections=KQueryDB.__QueryAllStockCollections()
		results={}
		tcountry=tcountry.upper()
		tmarket=tmarket.upper()
		for each_collection in allcollections:
			(country,market)=each_collection.split(".",2)
			if not (tcountry=='ALL' or country==tcountry):
				continue  #skip this market
			if not (tmarket=='ALL' or tmarket==market):
				continue  #skip this market
			cldailystocks=KQueryDB.Db[KDBUtility.GetCollectionName(country,market,True)]
			cursor=cldailystocks.find({},{"symbol":1,"latestinfo":1})
			for each_symbol in cursor:
				symbolname="%s.%s.%s" %(each_symbol['symbol'],market,country) 
				if 'latestinfo' in each_symbol:
					results[symbolname]=each_symbol['latestinfo']
				else:
					results[symbolname]=[]
		return results
	
	@staticmethod
	def QueryAllStocksStaticInfoBySymbolList(symbolist):	
		# query static stock information, 'static information' is all information which has only one value per stock 
		
		allcollections=KQueryDB.__QueryAllStockCollections()
		results=collections.OrderedDict()
		for symbol in symbolist:
			results[symbol]=[]
			
		#symbolist={'BAC.NYSE.US','AAPL.NASDAQ.US','BIDU.NASDAQ.US'}
		symgroup=KQueryDB.__ConvertSymbolList2DataFrameGroup(symbolist)
		for (country,market),sublist in symgroup:
			cldailystocks=KQueryDB.Db[KDBUtility.GetCollectionName(country,market,True)]
			actualist=sublist['Symbol'].tolist()
			cursor=cldailystocks.find({"symbol":{"$in":actualist}},{"symbol":1,"latestinfo":1})
			for each_symbol in cursor:
				symbolname="%s.%s.%s" %(each_symbol['symbol'],market,country) 
				if each_symbol==u'MHP.NYSE.US':
					test=True
				if 'latestinfo' in each_symbol:
					results[symbolname]=each_symbol['latestinfo']
				else:
					results[symbolname]=[]
					
		results=KQueryDB.__ConvertBasicInfoDict2DataFrame(results) # convert to dateframe
		return results
	
	@staticmethod
	def __ConvertBasicInfoDict2DataFrame(infodict):
		# convert basic info blocks into DataFrame (row will be stock name, column will be property)
		narray=np.empty([len(infodict),6])
		index=0
		symbolist=[]
		for each_symbol in infodict:
			if infodict[each_symbol]:
				narray[index,0:5]=infodict[each_symbol][0]
				narray[index,5]=infodict[each_symbol][1]
				symbolist.append(each_symbol)
				index+=1
			else:
				pass # do not add symbols without information
		narray.resize([len(symbolist),6])
		pdframe=pd.DataFrame(narray,index=symbolist,columns=['open','high','low','close','volumn','date'])
		return pdframe

	@staticmethod
	def __ConvertSymbolList2DataFrameGroup(symbolist):
		# convert symbollist composed of symbol.market.country into a structured pandas dateframe group
		rowcount=len(symbolist)
		
		newlist=[]
		for each_symbol in symbolist:
			sepsymbols=each_symbol.split('.')
			newlist.append(sepsymbols)
		#newlist=zip(*newlist)
		pdataframe=pd.DataFrame(newlist,columns=['Symbol','Market','Country'])
		grpdata=pdataframe.groupby(['Country','Market'])
		return grpdata
		
	@staticmethod
	def QueryStockLatestBasicsInfo(country,market,symbol):
		# get the lastest basic info field of the stock
		cldailystocks=KQueryDB.Db[KDBUtility.GetCollectionName(country,market,True)]
		cursor=cldailystocks.find_one({"symbol":symbol.upper()})	
		if cursor:
			return cursor['latestinfo']		
		else:
			return []	
		
	@staticmethod
	def QueryStockByDateRange(country, market,symbol):
		cldailystocks=KQueryDB.Db[KDBUtility.GetCollectionName("us",market,True)]
		cursor=cldailystocks.find_one({"symbol":symbol.upper()})
		#cursor=cldailystocks.find({"$and": [{"symbol":symbol.upper()},{"date":{"$gte":20130601}}]})#.sort("date") :can't apply since all dates live in the same record,can't return partial of the record
		if cursor:
			return [cursor['date'],cursor['price']]		
		else:
			return []